import numpy as np
import casadi
import meshcat.geometry as mg
import pinocchio as pin
from pinocchio import casadi as cpin  # 添加 casadi 的 pinocchio 绑定
from pinocchio.visualize import MeshcatVisualizer
import os
import sys
import time
import sys
import os
from utils.logger import run_time
from teleop.utils.logger import get_logger

logger = get_logger(__name__)

# 添加项目根目录到Python路径
project_root = os.path.abspath(os.path.join(os.path.dirname(__file__), '../..'))
if project_root not in sys.path:
    sys.path.append(project_root)

from teleop.utils.geo_math import rotation_matrix, euler_to_matrix

parent2_dir = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
print("parent2_dir:", parent2_dir)
sys.path.append(parent2_dir)

from teleop.utils.weighted_moving_filter import WeightedMovingFilter

# 设置URDF路径
urdf_path = os.path.join(parent2_dir, 'assets','arm_description9', 'urdf', 'arm_description9.urdf')
package_dir = os.path.join(parent2_dir, 'assets')

class FangNuoCarIK:
    def __init__(self, Unit_Test=False, Visualization=False):
        """
        初始化方法，加载机器人模型并设置参数
        
        参数:
            Unit_Test: 是否为单元测试模式
            Visualization: 是否启用可视化
        """
        np.set_printoptions(precision=4, suppress=True, linewidth=200)
        
        # 加载URDF模型
        try:
            self.robot = pin.RobotWrapper.BuildFromURDF(urdf_path, package_dir)
            print(f"成功加载URDF模型: {urdf_path}")
        except Exception as e:
            print(f"加载URDF模型失败: {e}")
            raise
        
        # 设置需要锁定的关节
        self.mixed_jointsToLockIDs = []
        
        # 构建简化机器人模型（移除锁定的关节）
        self.reduced_robot = self.robot
        if len(self.mixed_jointsToLockIDs) > 0:
            self.reduced_robot = self.robot.buildReducedRobot(list_of_joints_to_lock=self.mixed_jointsToLockIDs)
        
        # 添加末端执行器帧
        # 打印所有可用的帧名称，帮助调试
        # print("\n可用的帧名称:")
        # for i in range(self.reduced_robot.model.nframes):
        #     frame_name = self.reduced_robot.model.frames[i].name
        #     print(f"帧 {i}: {frame_name}")
        
        try:
            # 尝试获取左右手腕的帧ID，如果找不到，则使用其他可用的帧
            try:
                self.L_hand_id = self.reduced_robot.model.getFrameId("L_wrist_end")
                print(f"找到左手腕帧，ID: {self.L_hand_id}")
            except Exception as e:
                print(f"找不到'L_wrist_end'帧: {e}")
                # 尝试找到包含'left'或'arm'的帧
                for i in range(self.reduced_robot.model.nframes):
                    frame_name = self.reduced_robot.model.frames[i].name.lower()
                    if 'left' in frame_name or ('arm' in frame_name and 'left' in frame_name):
                        self.L_hand_id = i
                        print(f"使用替代左手帧: {self.reduced_robot.model.frames[i].name}, ID: {i}")
                        break
                else:
                    # 如果找不到合适的帧，使用第一个操作帧
                    self.L_hand_id = 1  # 通常第一个操作帧是基座之后的第一个帧
                    print(f"未找到合适的左手帧，使用默认帧 ID: {self.L_hand_id}")
            
            try:
                self.R_hand_id = self.reduced_robot.model.getFrameId("R_wrist_end")
                print(f"找到右手腕帧，ID: {self.R_hand_id}")
            except Exception as e:
                print(f"找不到'R_wrist_end'帧: {e}")
                # 尝试找到包含'right'或'arm'的帧
                for i in range(self.reduced_robot.model.nframes):
                    frame_name = self.reduced_robot.model.frames[i].name.lower()
                    if 'right' in frame_name or ('arm' in frame_name and 'right' in frame_name):
                        self.R_hand_id = i
                        print(f"使用替代右手帧: {self.reduced_robot.model.frames[i].name}, ID: {i}")
                        break
                else:
                    # 如果找不到合适的帧，使用与左手不同的帧
                    if hasattr(self, 'L_hand_id'):
                        self.R_hand_id = min(self.L_hand_id + 1, self.reduced_robot.model.nframes - 1)
                    else:
                        self.R_hand_id = 7  # 使用另一个默认帧
                    print(f"未找到合适的右手帧，使用默认帧 ID: {self.R_hand_id}")
            
            # 确保帧ID在有效范围内
            self.L_hand_id = min(self.L_hand_id, self.reduced_robot.model.nframes - 1)
            self.R_hand_id = min(self.R_hand_id, self.reduced_robot.model.nframes - 1)
            
            # 添加左右末端执行器帧
            try:
                self.L_ee_id = self.reduced_robot.model.addFrame(
                    pin.Frame("L_ee", self.L_hand_id, 0, pin.SE3.Identity(), pin.FrameType.OP_FRAME)
                )
                logger.info(f"成功添加左末端执行器帧，ID: {self.L_ee_id}")
            except Exception as e:
                logger.error(f"添加左末端执行器帧失败: {e}")
                self.L_ee_id = self.L_hand_id
                logger.info(f"使用左手腕帧代替，ID: {self.L_ee_id}")
            
            try:
                self.R_ee_id = self.reduced_robot.model.addFrame(
                    pin.Frame("R_ee", self.R_hand_id, 0, pin.SE3.Identity(), pin.FrameType.OP_FRAME)
                )
                logger.info(f"成功添加右末端执行器帧，ID: {self.R_ee_id}")
            except Exception as e:
                logger.error(f"添加右末端执行器帧失败: {e}")
                self.R_ee_id = self.R_hand_id
                logger.info(f"使用右手腕帧代替，ID: {self.R_ee_id}")
        except Exception as e:
            print(f"获取手腕帧ID失败: {e}")
            raise
        
        # 初始化数据和平滑滤波器
        self.init_data = np.zeros(self.reduced_robot.model.nq)
        self.smooth_filter = WeightedMovingFilter(np.array([0.4, 0.3, 0.2, 0.1]), 12)
        
        # 使用CasADi创建优化参数
        try:
            # 创建 Casadi 模型和数据，用于符号计算
            self.cmodel = cpin.Model(self.reduced_robot.model)
            self.cdata = self.cmodel.createData()
            
            # 创建符号变量
            self.cq = casadi.SX.sym("q", self.reduced_robot.model.nq, 1) 
            self.cTf_l = casadi.SX.sym("tf_l", 4, 4)
            self.cTf_r = casadi.SX.sym("tf_r", 4, 4)
            cpin.framesForwardKinematics(self.cmodel, self.cdata, self.cq)
            
            # 定义误差函数
            self.translational_error = casadi.Function(
                "translational_error",
                [self.cq, self.cTf_l, self.cTf_r],
                [
                    casadi.vertcat(
                        self.cdata.oMf[self.L_ee_id].translation - self.cTf_l[:3,3],
                        self.cdata.oMf[self.R_ee_id].translation - self.cTf_r[:3,3]
                    )
                ],
            )
            self.rotational_error = casadi.Function(
                "rotational_error",
                [self.cq, self.cTf_l, self.cTf_r],
                [
                    casadi.vertcat(
                        cpin.log3(self.cdata.oMf[self.L_ee_id].rotation @ self.cTf_l[:3,:3].T),
                        cpin.log3(self.cdata.oMf[self.R_ee_id].rotation @ self.cTf_r[:3,:3].T)
                    )
                ],
            )
            
            # 创建优化问题
            self.opti = casadi.Opti()
            
            # 创建优化变量（关节角度）
            self.var_q = self.opti.variable(self.reduced_robot.model.nq)
            
            # 创建参数（目标位姿）
            self.param_tf_l = self.opti.parameter(4, 4)  # 左手腕目标位姿
            self.param_tf_r = self.opti.parameter(4, 4)  # 右手腕目标位姿
            self.var_q_last = self.opti.parameter(self.reduced_robot.model.nq)  # 上一次的关节角度
            
            # 定义成本函数组件
            self.translational_cost = casadi.sumsqr(self.translational_error(self.var_q, self.param_tf_l, self.param_tf_r))
            self.rotation_cost = casadi.sumsqr(self.rotational_error(self.var_q, self.param_tf_l, self.param_tf_r))
            self.regularization_cost = casadi.sumsqr(self.var_q)
            self.smooth_cost = casadi.sumsqr(self.var_q - self.var_q_last)
            
            # 设置约束（关节角度上下限）
            self.opti.subject_to(self.opti.bounded(
                self.reduced_robot.model.lowerPositionLimit,
                self.var_q,
                self.reduced_robot.model.upperPositionLimit
            ))
            
            # 设置优化目标
            self.opti.minimize(50 * self.translational_cost + 0.5 * self.rotation_cost + 0.02 * self.regularization_cost + 0.1 * self.smooth_cost)
            
            # 设置求解器选项
            opts = {
                'ipopt': {
                    'print_level': 0,
                    'max_iter': 50,
                    'tol': 1e-6
                },
                'print_time': False,
                'calc_lam_p': False # https://github.com/casadi/casadi/wiki/FAQ:-Why-am-I-getting-%22NaN-detected%22in-my-optimization%3F
            }
            self.opti.solver('ipopt', opts)
            logger.info("成功初始化CasADi优化框架")
        except Exception as e:
            logger.error(f"初始化CasADi优化框架失败: {e}")
            logger.info("将使用原始的SciPy优化方法")
            self.use_casadi = False
        else:
            self.use_casadi = True
        
        # 可视化设置
        self.Visualization = Visualization
        self.vis = None
        
        if Visualization:
            logger.info("Visualization is enabled")
            
            try:
                # 尝试初始化可视化
                logger.info("正在初始化可视化...")
                self.vis = MeshcatVisualizer(self.reduced_robot.model, self.reduced_robot.collision_model, self.reduced_robot.visual_model)
                try:
                    self.vis.initViewer(open=True)
                    self.vis.loadViewerModel()
                    
                    # 启用末端执行器目标帧的显示
                    frame_viz_names = ['wrist_poses/left_wrist', 'wrist_poses/right_wrist']
                    FRAME_AXIS_POSITIONS = (
                        np.array([[0, 0, 0], [1, 0, 0],
                                  [0, 0, 0], [0, 1, 0],
                                  [0, 0, 0], [0, 0, 1]]).astype(np.float32).T
                    )
                    FRAME_AXIS_COLORS = (
                        np.array([[1, 0, 0], [1, 0.6, 0],
                                  [0, 1, 0], [0.6, 1, 0],
                                  [0, 0, 1], [0, 0.6, 1]]).astype(np.float32).T
                    )
                    axis_length = 0.2
                    axis_width = 20
                    for frame_viz_name in frame_viz_names:
                        self.vis.viewer[frame_viz_name].set_object(
                            mg.LineSegments(
                                mg.PointsGeometry(
                                    position=axis_length * FRAME_AXIS_POSITIONS,
                                    color=FRAME_AXIS_COLORS,
                                ),
                                mg.LineBasicMaterial(
                                    linewidth=axis_width,
                                    vertexColors=True,
                                ),
                            )
                        )
                        
                    logger.info("可视化初始化成功！")
            
                except Exception as e:
                    logger.error(f"可视化初始化失败: {e}")
                    logger.info("将继续运行但没有可视化功能")
                    self.Visualization = False
                logger.info("显示机器人模型...")
                self.vis.display(np.zeros(self.reduced_robot.model.nv))
            except Exception as e:
                logger.error(f"可视化模块加载失败: {e}")
                logger.info("将继续运行但没有可视化功能")
                self.Visualization = False
        else:
            logger.info("可视化功能已禁用")
    
    def scale_arms(self, human_left_pose, human_right_pose, human_arm_length=0.70, robot_arm_length=0.55):
        """
        缩放人类手臂位姿到机器人手臂尺寸
        
        参数:
            human_left_pose: 人类左手臂位姿矩阵 (4x4)
            human_right_pose: 人类右手臂位姿矩阵 (4x4)
            human_arm_length: 人类手臂长度 (默认0.60米)
            robot_arm_length: 机器人手臂长度 (默认0.75米)
            
        返回:
            robot_left_pose, robot_right_pose: 缩放后的机器人手臂位姿矩阵
        """
        scale_factor = robot_arm_length / human_arm_length
        robot_left_pose = human_left_pose.copy()
        robot_right_pose = human_right_pose.copy()
        robot_left_pose[:3, 3] *= scale_factor
        robot_right_pose[:3, 3] *= scale_factor
        return robot_left_pose, robot_right_pose
    @run_time
    def solve_ik(self, left_wrist, right_wrist, current_lr_arm_motor_q=None, current_lr_arm_motor_dq=None):
        """
        解决逆运动学问题，使用CasADi优化框架
        
        参数:
            left_wrist: 左手腕的目标位姿（齐次变换矩阵）
            right_wrist: 右手腕的目标位姿（齐次变换矩阵）
            current_lr_arm_motor_q: 当前关节角度
            current_lr_arm_motor_dq: 当前关节速度
            
        返回:
            关节角度和力矩，如果位置超出范围则返回(None, None)
        """
        # 初始化当前关节角度
        if current_lr_arm_motor_q is not None:
            self.init_data = current_lr_arm_motor_q
            
        # 设置初始猜测值
        self.opti.set_initial(self.var_q, self.init_data)
        
        # 校验手腕位置范围
        def check_pose_bounds(pose, name):
            x, y, z = pose[:3, 3]
            out_of_bounds = False
            
            if not (0.1 <= x <= 0.9):
                print(f"错误: {name}的x坐标{x:.3f}超出允许范围[0.1, 0.6]")
                out_of_bounds = True
            if not (-0.6 <= y <= 0.6):
                print(f"错误: {name}的y坐标{y:.3f}超出允许范围[-0.6, 0.6]")
                out_of_bounds = True
            if not (0.2 <= z <= 1.4):
                print(f"错误: {name}的z坐标{z:.3f}超出允许范围[0.2, 1.4]")
                out_of_bounds = True
                
            if out_of_bounds:
                return None
                
            return pose
        
        # 校验位置范围，如果超出范围则返回None
        left_wrist_checked = check_pose_bounds(left_wrist, "左手腕")
        right_wrist_checked = check_pose_bounds(right_wrist, "右手腕")
        
        # 如果任一位置超出范围，返回当前关节位置
        if left_wrist_checked is None or right_wrist_checked is None:
            print("错误: 手腕位置超出允许范围，保持当前关节位置")
            # 返回当前关节角度和零力矩
            current_q = self.init_data if current_lr_arm_motor_q is None else current_lr_arm_motor_q
            return current_q, np.zeros_like(current_q)
            
        left_wrist, right_wrist = left_wrist_checked, right_wrist_checked
              

        
        # 缩放手腕位姿
        left_wrist, right_wrist = self.scale_arms(left_wrist, right_wrist)
        
        # 可视化目标位姿
        # if self.Visualization and self.vis is not None:
        #     self.vis.viewer['L_ee_target'].set_transform(left_wrist)
        #     self.vis.viewer['R_ee_target'].set_transform(right_wrist)
        
        # 设置优化参数
        self.opti.set_value(self.param_tf_l, left_wrist)
        self.opti.set_value(self.param_tf_r, right_wrist)
        self.opti.set_value(self.var_q_last, self.init_data)  # 用于平滑
        
        try:
            # 求解优化问题
            sol = self.opti.solve()
            # sol = self.opti.solve_limited()
            
            # 获取优化结果
            sol_q = self.opti.value(self.var_q)
            
            # 平滑关节角度
            self.smooth_filter.add_data(sol_q)
            sol_q = self.smooth_filter.filtered_data
            
            # 计算关节速度
            if current_lr_arm_motor_dq is not None:
                v = current_lr_arm_motor_dq * 0.0
            else:
                v = (sol_q - self.init_data) * 0.0
            
            # 更新初始数据
            self.init_data = sol_q
            
            # 计算关节力矩
            # sol_tauff = pin.rnea(self.reduced_robot.model, self.reduced_robot.data, sol_q, v, np.zeros(self.reduced_robot.model.nv))
            
            # 可视化结果
            if self.Visualization and self.vis is not None:
                self.vis.display(sol_q)  # for visualization
            
            # 输出结果信息
            # print(f"左手腕目标位置: {left_wrist[:3, 3]}")
            # print(f"右手腕目标位置: {right_wrist[:3, 3]}")
            # print(f"左手腕关节角度: {sol_q[:6]}")
            # print(f"右手腕关节角度: {sol_q[6:12]}")

            # print(f"优化成功完成")
            
            return sol_q, np.zeros(self.reduced_robot.model.nv)
        
        except Exception as e:
            print(f"优化失败，直接返回当前关节位置: {e}")
            
            # 如果求解失败，直接返回当前的关节位置和零力矩
            if self.Visualization and self.vis is not None and current_lr_arm_motor_q is not None:
                self.vis.display(current_lr_arm_motor_q)  # 可视化当前关节位置
            
            print(f"返回当前关节位置: {current_lr_arm_motor_q}")
            return current_lr_arm_motor_q, np.zeros(self.reduced_robot.model.nv)

    def get_frame_pose(self, q, frame_id=None):
        """
        获取指定坐标系的任务空间位姿
        
        参数:
            q: 关节角度配置
            frame_id: 要获取位姿的坐标系ID，如果为None，则返回左右手腕的位姿
            
        返回:
            如果frame_id为None，返回(left_pose, right_pose)，否则返回指定frame_id的位姿
        """
        # 计算正向运动学
        pin.forwardKinematics(self.reduced_robot.model, self.reduced_robot.data, q)
        pin.updateFramePlacements(self.reduced_robot.model, self.reduced_robot.data)
        
        # 如果指定了特定的frame_id，返回该坐标系的位姿
        if frame_id is not None:
            # 确保frame_id在有效范围内
            if frame_id >= len(self.reduced_robot.data.oMf):
                print(f"警告: 指定的frame_id {frame_id} 超出范围，最大值为 {len(self.reduced_robot.data.oMf)-1}")
                frame_id = min(frame_id, len(self.reduced_robot.data.oMf)-1)
            
            # 获取位姿矩阵
            pose = self.reduced_robot.data.oMf[frame_id]
            # 转换为4x4齐次变换矩阵
            pose_matrix = np.eye(4)
            pose_matrix[:3, :3] = pose.rotation
            pose_matrix[:3, 3] = pose.translation
            
            return pose_matrix
        
        # 否则返回左右手腕的位姿
        # 确保帧ID在有效范围内
        L_hand_id = min(self.L_hand_id, len(self.reduced_robot.data.oMf)-1)
        R_hand_id = min(self.R_hand_id, len(self.reduced_robot.data.oMf)-1)
        
        # 获取左手腕位姿
        left_pose = self.reduced_robot.data.oMf[L_hand_id]
        left_matrix = np.eye(4)
        left_matrix[:3, :3] = left_pose.rotation
        left_matrix[:3, 3] = left_pose.translation
        
        # 获取右手腕位姿
        right_pose = self.reduced_robot.data.oMf[R_hand_id]
        right_matrix = np.eye(4)
        right_matrix[:3, :3] = right_pose.rotation
        right_matrix[:3, 3] = right_pose.translation
        
        # 如果可视化已启用，显示手腕位姿
        if self.Visualization and self.vis is not None:
            self.display_wrist_poses(left_matrix, right_matrix)
        
        return left_matrix, right_matrix
        
    def display_wrist_poses(self, left_wrist_matrix, right_wrist_matrix, axis_length=0.1, axis_width=5):
        """在meshcat中显示左右手腕的位姿
        
        参数:
            left_wrist_matrix: 左手腕位姿矩阵 (4x4)
            right_wrist_matrix: 右手腕位姿矩阵 (4x4)
            axis_length: 坐标轴长度
            axis_width: 坐标轴宽度
        """
        if self.vis is None or not self.Visualization:
            return
            
        import meshcat.geometry as mg
        
        # 定义坐标轴的位置和颜色
        FRAME_AXIS_POSITIONS = (
            np.array([[0, 0, 0], [1, 0, 0],
                      [0, 0, 0], [0, 1, 0],
                      [0, 0, 0], [0, 0, 1]]).astype(np.float32).T
        )
        FRAME_AXIS_COLORS = (
            np.array([[1, 0, 0], [1, 0.6, 0],
                      [0, 1, 0], [0.6, 1, 0],
                      [0, 0, 1], [0, 0.6, 1]]).astype(np.float32).T
        )
        
        # 创建左手腕位姿可视化
        left_wrist_viz_name = "wrist_poses/left_wrist"
        # 直接创建坐标轴对象，不需要检查是否存在
        try:
            self.vis.viewer[left_wrist_viz_name].set_object(
                mg.LineSegments(
                    mg.PointsGeometry(
                        position=axis_length * FRAME_AXIS_POSITIONS,
                        color=FRAME_AXIS_COLORS,
                    ),
                    mg.LineBasicMaterial(
                        linewidth=axis_width,
                        vertexColors=True,
                    ),
                )
            )
        except Exception as e:
            print(f"创建左手腕可视化失败: {e}")
        
        # 更新左手腕位姿
        try:
            self.vis.viewer[left_wrist_viz_name].set_transform(left_wrist_matrix)
        except Exception as e:
            print(f"更新左手腕位姿失败: {e}")
        
        # 创建右手腕位姿可视化
        right_wrist_viz_name = "wrist_poses/right_wrist"
        try:
            self.vis.viewer[right_wrist_viz_name].set_object(
                mg.LineSegments(
                    mg.PointsGeometry(
                        position=axis_length * FRAME_AXIS_POSITIONS,
                        color=FRAME_AXIS_COLORS,
                    ),
                    mg.LineBasicMaterial(
                        linewidth=axis_width,
                        vertexColors=True,
                    ),
                )
            )
        except Exception as e:
            print(f"创建右手腕可视化失败: {e}")
            
        # 更新右手腕位姿
        try:
            self.vis.viewer[right_wrist_viz_name].set_transform(right_wrist_matrix)
        except Exception as e:
            print(f"更新右手腕位姿失败: {e}")

def rotation_matrix(axis, theta):
    """
    计算绕指定轴旋转指定角度的旋转矩阵
    
    参数:
        axis: 旋转轴向量，需要是单位向量 [x, y, z]
        theta: 旋转角度（弧度）
    
    返回:
        3x3 旋转矩阵
    """
    # 确保轴是单位向量
    axis = np.array(axis, dtype=np.float64)
    axis = axis / np.linalg.norm(axis)
    
    # 使用罗德里格斯旋转公式 (Rodrigues' rotation formula)
    a = np.cos(theta / 2.0)
    b, c, d = -axis * np.sin(theta / 2.0)
    
    # 四元数元素
    aa, bb, cc, dd = a * a, b * b, c * c, d * d
    bc, ad, ac, ab, bd, cd = b * c, a * d, a * c, a * b, b * d, c * d
    
    # 构建旋转矩阵
    rotation_matrix = np.array([
        [aa + bb - cc - dd, 2 * (bc + ad), 2 * (bd - ac)],
        [2 * (bc - ad), aa + cc - bb - dd, 2 * (cd + ab)],
        [2 * (bd + ac), 2 * (cd - ab), aa + dd - bb - cc]
    ])
    
    return rotation_matrix

# 单元测试代码
if __name__ == "__main__":
    # 创建IK求解器实例
    ik_solver = FangNuoCarIK(Unit_Test=True, Visualization=True)

# 左手腕位姿: 
# [[ 0.0872  0.      0.9962  0.2679]
#  [-0.      1.     -0.      0.25  ]
#  [-0.9962 -0.      0.0872  0.7831]
#  [ 0.      0.      0.      1.    ]]
# 右手腕位姿: 
# [[ 0.0872 -0.      0.9962  0.2679]
#  [ 0.      1.      0.     -0.25  ]
#  [-0.9962  0.      0.0872  0.7831]
#  [ 0.      0.      0.      1.    ]]
    left_wrist_pose, right_wrist_pose = ik_solver.get_frame_pose(ik_solver.init_data)
    print("左手腕位姿:", left_wrist_pose)
    print("右手腕位姿:", right_wrist_pose)
    ik_solver.solve_ik(left_wrist_pose, right_wrist_pose)
    
    ik_solver.vis.display(ik_solver.init_data)
    
    input("按回车键退出...")
    
    # 设置测试目标位姿
    left_target = np.eye(4)
    # 创建绕Y轴旋转-90度的旋转矩阵
    # 使用rotation_matrix函数计算旋转矩阵
    # left_target[:3, :3] = rotation_matrix([0, 1, 0], np.pi/2)
    # left_target = left_target @ euler_to_matrix([0, 90, 0], sequence='xyz')
    # left_target[:3, 3] = [0.23, 0.25, 0.79]  # 左手腕目标位置
    
    # right_target = np.eye(4)
    # # right_target[:3, :3] = rotation_matrix([0, 1, 0], np.pi/2)
    # right_target = right_target @ euler_to_matrix([0, 90, 0], sequence='xyz')
    # right_target[:3, 3] = [0.23, -0.25, 0.79]  # 右手腕目标位置
    
    # 求解IK
    # q, tau = ik_solver.solve_ik(left_target, right_target)
    
    # print("\n最终关节角度:")
    # print(q)
    # print("\n关节力矩:")
    # print(tau)
    # print("hand pose:", ik_solver.get_frame_pose(q))
    # time.sleep(3)
    # ik_solver.vis.display(np.zeros(ik_solver.reduced_robot.model.nv))
    # time.sleep(3)
    # ik_solver.vis.display(q)
    

    # input("继续...")
    # left_target[:3, 3] = [0.5, 0.25, 1]
    # right_target[:3, 3] = [0.5, -0.25, 1]
    # q, tau = ik_solver.solve_ik(left_target, right_target)
    # print("\n最终关节角度:")
    # print(q)
    # print("\n关节力矩:")
    # print(tau)
    # print("hand pose:", ik_solver.get_frame_pose(q))
    
    # # 如果启用了可视化，等待用户查看
    # if ik_solver.Visualization:
    #     input("按回车键退出...")

    # print(rotation_matrix([0, 1, 0], np.pi/2))